"""
RAG (Retrieval Augmented Generation) 管理器

提供检索增强生成的核心功能，包括文档索引、向量检索、问答生成等。
基于现有的LLM管理器、嵌入模型管理器和向量存储管理器构建。
"""

import os
import logging
from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime

from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain.chains import RetrievalQA
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain

from core.llm_manager import llm_manager
from core.embedding_manager import embedding_manager
from core.vector_store import vector_store_manager
from utils.document_loader import DocumentLoader
from utils.text_splitter import SmartTextSplitter
from utils.helpers import format_docs, clean_text, filter_documents_by_length, deduplicate_documents

logger = logging.getLogger(__name__)

class RAGManager:
    """RAG管理器 - 基于现有组件构建的检索增强生成系统"""
    
    def __init__(self, 
                 chunk_size: int = 1000,
                 chunk_overlap: int = 200,
                 top_k: int = 5):
        """
        初始化RAG管理器
        
        Args:
            chunk_size (int): 文本块大小
            chunk_overlap (int): 文本块重叠大小
            top_k (int): 检索文档数量
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.top_k = top_k
        
        # 使用现有的管理器
        self.llm_manager = llm_manager
        self.embedding_manager = embedding_manager
        # 导入全局向量存储管理器实例
        self.vector_store_manager = vector_store_manager
        
        # 文档处理组件
        self.document_loader = DocumentLoader()
        self.text_splitter = SmartTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap
        )
        
        # QA链
        self.qa_chain = None
        
        # 文档统计
        self.doc_stats = {
            'total_documents': 0,
            'total_chunks': 0,
            'last_updated': None,
            'sources': []
        }
        
        self._initialize_qa_chain()
    
    def _initialize_qa_chain(self):
        """初始化问答链"""
        try:
            # 检查依赖组件是否就绪
            if not self.llm_manager.get_current_model():
                logger.warning("LLM未就绪，问答链将在LLM就绪后创建")
                return
            
            if not self.vector_store_manager.is_available():
                logger.warning("向量存储未就绪，问答链将在向量存储就绪后创建")
                return
            
            self._create_qa_chain()
            logger.info("RAG管理器初始化成功")
            
        except Exception as e:
            logger.error(f"RAG管理器初始化失败: {e}")
    
    def _create_qa_chain(self):
        """创建问答链"""
        try:
            # 获取向量存储的检索器
            if not self.vector_store_manager.vector_store:
                logger.warning("向量存储不可用，无法创建问答链")
                return
            
            retriever = self.vector_store_manager.vector_store.as_retriever(
                search_type="similarity",
                search_kwargs={"k": self.top_k}
            )
            
            # 创建提示模板
            system_template = """你是一个专业的AI助手，基于给定的文档内容回答用户问题。

指导原则：
1. 仅基于提供的文档内容回答问题
2. 如果文档中没有相关信息，明确告知用户
3. 提供准确、详细且有帮助的回答
4. 在适当的地方引用文档来源
5. 保持回答的逻辑性和条理性

文档内容：
{context}

用户问题：{input}

请基于以上文档内容回答用户问题："""

            prompt = ChatPromptTemplate.from_template(system_template)
            
            # 获取当前LLM
            llm = self.llm_manager.get_current_model()
            if not llm:
                logger.error("无法获取LLM实例")
                return
            
            # 创建文档链
            document_chain = create_stuff_documents_chain(llm, prompt)
            
            # 创建检索链
            self.qa_chain = create_retrieval_chain(retriever, document_chain)
            
            logger.info("问答链创建成功")
            
        except Exception as e:
            logger.error(f"创建问答链失败: {e}")
    
    def add_documents_from_file(self, file_path: str) -> Dict[str, Any]:
        """
        从文件添加文档
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            Dict[str, Any]: 处理结果
        """
        try:
            # 加载文档
            docs = self.document_loader.load_file(file_path)
            
            if not docs:
                return {
                    'success': False,
                    'message': f'无法从文件 {file_path} 加载文档',
                    'documents_added': 0
                }
            
            return self._process_and_add_documents(docs, file_path)
            
        except Exception as e:
            logger.error(f"从文件添加文档失败: {e}")
            return {
                'success': False,
                'message': f'添加文档失败: {e}',
                'documents_added': 0
            }
    
    def add_documents_from_directory(self, directory_path: str) -> Dict[str, Any]:
        """
        从目录添加文档
        
        Args:
            directory_path (str): 目录路径
            
        Returns:
            Dict[str, Any]: 处理结果
        """
        try:
            # 加载目录中的所有文档
            docs = self.document_loader.load_directory(directory_path)
            
            if not docs:
                return {
                    'success': False,
                    'message': f'目录 {directory_path} 中没有找到支持的文档',
                    'documents_added': 0
                }
            
            return self._process_and_add_documents(docs, directory_path)
            
        except Exception as e:
            logger.error(f"从目录添加文档失败: {e}")
            return {
                'success': False,
                'message': f'添加文档失败: {e}',
                'documents_added': 0
            }
    
    def add_documents_from_urls(self, urls: List[str]) -> Dict[str, Any]:
        """
        从URL添加文档
        
        Args:
            urls (List[str]): URL列表
            
        Returns:
            Dict[str, Any]: 处理结果
        """
        try:
            # 加载URL文档
            docs = self.document_loader.load_urls(urls)
            
            if not docs:
                return {
                    'success': False,
                    'message': '无法从提供的URL加载文档',
                    'documents_added': 0
                }
            
            return self._process_and_add_documents(docs, f"URLs: {', '.join(urls)}")
            
        except Exception as e:
            logger.error(f"从URL添加文档失败: {e}")
            return {
                'success': False,
                'message': f'添加文档失败: {e}',
                'documents_added': 0
            }
    
    def add_text(self, text: str, metadata: Dict[str, Any] = None) -> Dict[str, Any]:
        """
        直接添加文本内容到知识库
        
        Args:
            text (str): 要添加的文本内容
            metadata (Dict[str, Any], optional): 文档元数据
            
        Returns:
            Dict[str, Any]: 处理结果
        """
        try:
            if not text or not text.strip():
                return {
                    'success': False,
                    'message': '文本内容不能为空',
                    'documents_added': 0
                }
            
            # 创建文档对象
            if metadata is None:
                metadata = {}
            
            # 添加默认元数据
            metadata.setdefault('source', 'direct_text')
            metadata.setdefault('type', 'text')
            
            doc = Document(page_content=text.strip(), metadata=metadata)
            docs = [doc]
            
            return self._process_and_add_documents(docs, f"直接文本: {metadata.get('source', 'unknown')}")
            
        except Exception as e:
            logger.error(f"添加文本失败: {e}")
            return {
                'success': False,
                'message': f'添加文本失败: {e}',
                'documents_added': 0
            }
    
    def _process_and_add_documents(self, docs: List[Document], source: str) -> Dict[str, Any]:
        """
        处理并添加文档到向量数据库
        
        Args:
            docs (List[Document]): 文档列表
            source (str): 来源描述
            
        Returns:
            Dict[str, Any]: 处理结果
        """
        try:
            # 检查向量存储是否可用
            if not self.vector_store_manager.is_available():
                return {
                    'success': False,
                    'message': '向量存储不可用，请检查嵌入模型配置',
                    'documents_added': 0
                }
            
            # 文档预处理
            filtered_docs = filter_documents_by_length(docs, min_length=50)
            deduplicated_docs = deduplicate_documents(filtered_docs)
            
            # 分割文档
            chunks = self.text_splitter.split_documents(deduplicated_docs)
            
            if not chunks:
                return {
                    'success': False,
                    'message': '文档处理后没有生成有效的文本块',
                    'documents_added': 0
                }
            
            # 添加到向量数据库
            success = self.vector_store_manager.add_documents(chunks)
            
            if not success:
                return {
                    'success': False,
                    'message': '向量存储添加文档失败',
                    'documents_added': 0
                }
            
            # 重新创建QA链（因为向量存储可能刚刚初始化）
            if not self.qa_chain:
                self._create_qa_chain()
            
            # 更新统计信息
            self.doc_stats['total_documents'] += len(deduplicated_docs)
            self.doc_stats['total_chunks'] += len(chunks)
            self.doc_stats['last_updated'] = datetime.now().isoformat()
            
            if source not in self.doc_stats['sources']:
                self.doc_stats['sources'].append(source)
            
            return {
                'success': True,
                'message': f'成功添加 {len(deduplicated_docs)} 个文档，生成 {len(chunks)} 个文本块',
                'documents_added': len(deduplicated_docs),
                'chunks_generated': len(chunks),
                'source': source
            }
            
        except Exception as e:
            logger.error(f"处理文档失败: {e}")
            return {
                'success': False,
                'message': f'处理文档失败: {e}',
                'documents_added': 0
            }
    
    def ask_question(self, question: str) -> Dict[str, Any]:
        """
        问答
        
        Args:
            question (str): 用户问题
            
        Returns:
            Dict[str, Any]: 回答结果
        """
        # 检查组件状态
        if not self.llm_manager.get_current_model():
            return {
                'success': False,
                'answer': 'LLM未就绪，请先配置语言模型',
                'sources': []
            }
        
        if not self.qa_chain:
            return {
                'success': False,
                'answer': '知识库为空，请先添加文档',
                'sources': []
            }
        
        try:
            # 执行问答
            result = self.qa_chain.invoke({"input": question})
            
            # 提取来源信息
            sources = []
            if 'context' in result:
                for doc in result['context']:
                    source_info = {
                        'source': doc.metadata.get('source', 'Unknown'),
                        'content_preview': doc.page_content[:200] + '...' if len(doc.page_content) > 200 else doc.page_content
                    }
                    sources.append(source_info)
            
            return {
                'success': True,
                'answer': result.get('answer', '无法生成回答'),
                'sources': sources,
                'question': question
            }
            
        except Exception as e:
            logger.error(f"问答失败: {e}")
            return {
                'success': False,
                'answer': f'问答过程中发生错误: {e}',
                'sources': []
            }
    
    def search_documents(self, query: str, k: int = None) -> List[Dict[str, Any]]:
        """
        搜索相关文档
        
        Args:
            query (str): 搜索查询
            k (int): 返回文档数量
            
        Returns:
            List[Dict[str, Any]]: 搜索结果
        """
        if not self.vector_store_manager.is_available():
            return []
        
        try:
            search_k = k or self.top_k
            docs = self.vector_store_manager.similarity_search(query, k=search_k)
            
            results = []
            for i, doc in enumerate(docs, 1):
                results.append({
                    'rank': i,
                    'content': doc.page_content,
                    'source': doc.metadata.get('source', 'Unknown'),
                    'metadata': doc.metadata
                })
            
            return results
            
        except Exception as e:
            logger.error(f"搜索文档失败: {e}")
            return []
    
    def get_knowledge_base_info(self) -> Dict[str, Any]:
        """
        获取知识库信息
        
        Returns:
            Dict[str, Any]: 知识库信息
        """
        vector_info = self.vector_store_manager.get_collection_info()
        
        info = {
            'has_vectorstore': self.vector_store_manager.is_available(),
            'vector_count': vector_info.get('collection_count', 0),
            'statistics': self.doc_stats.copy(),
            'llm_ready': self.llm_manager.get_current_model() is not None,
            'embedding_ready': self.embedding_manager.is_ready()
        }
        
        return info
    
    def clear_knowledge_base(self) -> bool:
        """
        清空知识库
        
        Returns:
            bool: 是否成功
        """
        try:
            # 清空向量存储
            success = self.vector_store_manager.clear_collection()
            
            if success:
                # 重置QA链
                self.qa_chain = None
                
                # 重置统计信息
                self.doc_stats = {
                    'total_documents': 0,
                    'total_chunks': 0,
                    'last_updated': None,
                    'sources': []
                }
                
                logger.info("知识库已清空")
                return True
            else:
                logger.error("清空知识库失败")
                return False
            
        except Exception as e:
            logger.error(f"清空知识库失败: {e}")
            return False
    
    def export_knowledge_base_summary(self) -> str:
        """
        导出知识库摘要
        
        Returns:
            str: 知识库摘要
        """
        info = self.get_knowledge_base_info()
        llm_info = self.llm_manager.get_model_info()
        embedding_info = self.embedding_manager.get_model_info()
        
        summary = f"""
知识库摘要报告
================

系统状态：
- LLM状态：{'就绪' if info['llm_ready'] else '未就绪'}
- 嵌入模型状态：{'就绪' if info['embedding_ready'] else '未就绪'}
- 向量存储状态：{'就绪' if info['has_vectorstore'] else '未就绪'}

知识库信息：
- 向量数量：{info['vector_count']}
- 文档数量：{info['statistics']['total_documents']}
- 文本块数量：{info['statistics']['total_chunks']}
- 最后更新：{info['statistics']['last_updated'] or '未知'}

数据源：
{chr(10).join(f"- {source}" for source in info['statistics']['sources']) if info['statistics']['sources'] else "- 无"}

模型配置：
- 语言模型：{llm_info.get('model_name', '未配置')}
- 嵌入模型：{embedding_info.get('model_name', '未配置')}
- 文本块大小：{self.chunk_size}
- 文本块重叠：{self.chunk_overlap}
- 检索文档数：{self.top_k}
"""
        
        return summary.strip()
    
    def is_ready(self) -> bool:
        """
        检查RAG系统是否就绪
        
        Returns:
            bool: 是否就绪
        """
        return (
            self.llm_manager.get_current_model() is not None and
            self.embedding_manager.is_ready() and
            self.vector_store_manager.is_available()
        ) 